﻿using System;
using System.Drawing;
using System.Linq;
using System.Runtime.InteropServices;
using System.Text;

namespace AssociationManager
{
    public class IconManager
    {
        public enum IconSize { Small, Large }

        public int Count { get; private set; }

        private string _libaryPath;
        private ApplicationIcon _applicationIcon;

        public IconManager(ApplicationIcon icon)
        {
            _applicationIcon = icon;
            if (_applicationIcon == null)
                return;

            _libaryPath = icon.IconLibraryPath;

            if (!string.IsNullOrEmpty(_libaryPath))
            {
                Count = (int)Shell32.ExtractIconEx(_libaryPath, -1, null, null, 0);
            }
        }

        public IconManager(string libaryPath)
        {
            _libaryPath = libaryPath;

            Count = (int)Shell32.ExtractIconEx(_libaryPath, -1, null, null, 0);
        }

        public Icon GetIcon(IconSize iconSize, string filename)
        {
            if (!string.IsNullOrEmpty(_applicationIcon?.IExtractIconGuid) && !string.IsNullOrEmpty(filename))
            {
                Shell32.SHFILEINFO shfi = new Shell32.SHFILEINFO();
                try
                {
                    uint flags = Shell32.SHGFI_ICON | Shell32.SHGFI_USEFILEATTRIBUTES;
                    switch (iconSize)
                    {
                        case IconSize.Small: flags = flags | Shell32.SHGFI_SMALLICON; break;
                        case IconSize.Large: flags = flags | Shell32.SHGFI_LARGEICON; break;
                    }

                    Shell32.SHGetFileInfo(filename, 0, ref shfi, (uint)Marshal.SizeOf(shfi), flags);

                    return (Icon)Icon.FromHandle(shfi.hIcon).Clone();
                }
                catch (Exception ex)
                {
                    throw new Exception("Failed to extract icon", ex);
                }
                finally
                {
                    try
                    {
                        if (shfi.hIcon != null)
                            Shell32.DestroyIcon(shfi.hIcon);
                    }
                    catch { }
                }
            }

            return GetIcon(iconSize, _applicationIcon?.IconIndex);
        }

        public Icon GetIcon(IconSize iconSize, int? iconIndex)
        {
            if (!iconIndex.HasValue)
                iconIndex = 0;

            uint readIconCount = 0;
            IntPtr[] hDummy = new IntPtr[1] { IntPtr.Zero };
            IntPtr[] hIconEx = new IntPtr[1] { IntPtr.Zero };
            try
            {
                switch (iconSize)
                {
                    case IconSize.Small:
                        {
                            readIconCount = Shell32.ExtractIconEx(_libaryPath, iconIndex.Value, hDummy, hIconEx, 1);
                            break;
                        }
                    case IconSize.Large:
                        {
                            readIconCount = Shell32.ExtractIconEx(_libaryPath, iconIndex.Value, hIconEx, hDummy, 1);
                            break;
                        }
                }

                if (readIconCount > 0 && hIconEx[0] != IntPtr.Zero)
                    return (Icon)Icon.FromHandle(hIconEx[0]).Clone();
            }
            catch (Exception ex)
            {
                throw new Exception("Failed to extract icon", ex);
            }
            finally
            {
                foreach (IntPtr ptr in hIconEx)
                {
                    if (ptr == IntPtr.Zero)
                        continue;

                    try
                    {
                        Shell32.DestroyIcon(ptr);
                    }
                    catch { }
                }

                foreach (IntPtr ptr in hDummy)
                {
                    if (ptr == IntPtr.Zero)
                        continue;

                    try
                    {
                        Shell32.DestroyIcon(ptr);
                    }
                    catch { }
                }
            }

            return null;
        }

        private static class Shell32
        {
            public const uint SHGFI_ICON = 0x000000100;     // get icon
            public const uint SHGFI_DISPLAYNAME = 0x000000200;     // get display name
            public const uint SHGFI_TYPENAME = 0x000000400;     // get type name
            public const uint SHGFI_ATTRIBUTES = 0x000000800;     // get attributes
            public const uint SHGFI_ICONLOCATION = 0x000001000;     // get icon location
            public const uint SHGFI_EXETYPE = 0x000002000;     // return exe type
            public const uint SHGFI_SYSICONINDEX = 0x000004000;     // get system icon index
            public const uint SHGFI_LINKOVERLAY = 0x000008000;     // put a link overlay on icon
            public const uint SHGFI_SELECTED = 0x000010000;     // show icon in selected state
            public const uint SHGFI_ATTR_SPECIFIED = 0x000020000;     // get only specified attributes
            public const uint SHGFI_LARGEICON = 0x000000000;     // get large icon
            public const uint SHGFI_SMALLICON = 0x000000001;     // get small icon
            public const uint SHGFI_OPENICON = 0x000000002;     // get open icon
            public const uint SHGFI_SHELLICONSIZE = 0x000000004;     // get shell size icon
            public const uint SHGFI_PIDL = 0x000000008;     // pszPath is a pidl
            public const uint SHGFI_USEFILEATTRIBUTES = 0x000000010;     // use passed dwFileAttribute
            public const uint SHGFI_ADDOVERLAYS = 0x000000020;     // apply the appropriate overlays
            public const uint SHGFI_OVERLAYINDEX = 0x000000040;     // Get the index of the overlay

            public const uint FILE_ATTRIBUTE_DIRECTORY = 0x00000010;
            public const uint FILE_ATTRIBUTE_NORMAL = 0x00000080;

            private const int MAX_PATH = 260;

            [DllImport("shell32.dll", CharSet = CharSet.Auto)]
            public static extern uint ExtractIconEx(string szFileName, int nIconIndex, IntPtr[] phiconLarge, IntPtr[] phiconSmall, uint nIcons);

            [DllImport("user32.dll", EntryPoint = "DestroyIcon", SetLastError = true)]
            public static extern int DestroyIcon(IntPtr hIcon);

            [StructLayout(LayoutKind.Sequential)]
            public struct SHFILEINFO
            {
                public const int NAMESIZE = 80;
                public IntPtr hIcon;
                public int iIcon;
                public uint dwAttributes;
                [MarshalAs(UnmanagedType.ByValTStr, SizeConst = MAX_PATH)]
                public string szDisplayName;
                [MarshalAs(UnmanagedType.ByValTStr, SizeConst = NAMESIZE)]
                public string szTypeName;
            };

            [DllImport("Shell32.dll")]
            public static extern IntPtr SHGetFileInfo(string pszPath, uint dwFileAttributes, ref SHFILEINFO psfi, uint cbFileInfo, uint uFlags);
        }
    }
}
